package com.jiangsong.utils;

import cn.hutool.core.bean.BeanUtil;
import com.jiangsong.utils.base.MultiplyQueryPage;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Function;

public class MultiplyQueryUtil<I extends MultiplyQueryPage, R> {

    private Function<I, List<R>> queryFunc;

    private int multiplyCount = 0;

    private int pageSize = 5000;

    private int queryLimit = 50000;

    public static <I extends MultiplyQueryPage, R> MultiplyQueryUtil<I, R> of(Function<I, List<R>> function) {
        MultiplyQueryUtil<I, R> irMultiplyQueryUtil = new MultiplyQueryUtil<>();
        irMultiplyQueryUtil.queryFunc = function;
        return irMultiplyQueryUtil;
    }

    public static <I extends MultiplyQueryPage, R> MultiplyQueryUtil<I, R> of(int multiplyCount, Function<I, List<R>> function) {
        MultiplyQueryUtil<I, R> irMultiplyQueryUtil = new MultiplyQueryUtil<>();
        irMultiplyQueryUtil.queryFunc = function;
        irMultiplyQueryUtil.multiplyCount = multiplyCount;
        return irMultiplyQueryUtil;
    }

    public MultiplyQueryUtil<I, R> initPageSize(int pageSize) {
        this.pageSize = pageSize;
        return this;
    }

    public MultiplyQueryUtil<I, R> initQueryLimit(int queryLimit) {
        this.queryLimit = queryLimit;
        return this;
    }

    /**
     * 串行分批次查询
     *
     * @param queryInput
     * @return
     */
    public List<R> execute(I queryInput) {
        Assert.notNull(queryInput, "queryInput不能为空");
        int maxRound = (queryLimit - 1) / pageSize + 1;
        int start = 0;
        List<R> res = new ArrayList<>();

        for (int i = 0; i < maxRound; i++) {
            queryInput.setStartIndex(start);
            queryInput.setPageSize(pageSize);
            start += pageSize;
            List<R> partAns = queryFunc.apply(queryInput);
            if (null == partAns) {
                // 查询结束
                break;
            }
            res.addAll(partAns);
            if (partAns.size() < pageSize) {
                // 查询结束
                break;
            }
        }
        return res;
    }


    public List<R> execute(I queryInput, Function<I, Integer> countFunc) {
        Assert.notNull(queryInput, "queryInput不能为空");
        Assert.isTrue(multiplyCount > 0, "创建对象时未指定并发数");
        Integer maxCount = countFunc.apply(queryInput);
        return execute(queryInput, maxCount);
    }

    public List<R> execute(I queryInput, int maxCount) {
        Assert.notNull(queryInput, "queryInput不能为空");
        Assert.isTrue(multiplyCount > 0, "创建对象时未指定并发数");
        if (maxCount <= 0) {
            return new ArrayList<>();
        }
        Assert.isTrue(maxCount <= queryLimit, "数量超过上限" + queryLimit);
        if (maxCount <= pageSize) {
            System.out.println("maxCount <= pageSize" + pageSize + " " + maxCount);
            queryInput.setStartIndex(0);
            queryInput.setPageSize(maxCount);
            return queryFunc.apply(queryInput);
        }
        int queryRound = (maxCount - 1) / pageSize + 1;

        ExecutorService executorService = Executors.newFixedThreadPool(Integer.min(queryRound, multiplyCount));
        CompletableFuture<?>[] completableFutures = new CompletableFuture[queryRound];

        int start = 0;
        for (int i = 0; i < queryRound; i++) {
            // 复制对象
            Class<I> aClass = (Class<I>) queryInput.getClass();
            I queryPage = BeanUtil.copyProperties(queryInput, aClass);
            queryPage.setStartIndex(start);
            queryPage.setPageSize(pageSize);
            start += pageSize;
            completableFutures[i] = CompletableFuture.supplyAsync(() -> queryFunc.apply(queryPage), executorService);
        }

        List<R> res = new ArrayList<>();

        CompletableFuture<?> all = CompletableFuture.allOf(completableFutures);
        try {
            all.get();
            for (CompletableFuture<?> each : completableFutures) {

                res.addAll((Collection<? extends R>) each.get());
            }
        } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
        System.out.println("销毁线程池");
        executorService.shutdown();

        return res;

    }


}
